'''
Created on Mar 22, 2011

@author: oabalbin
'''

import os
import subprocess
from optparse import OptionParser
from exome.jobs.job_runner import qsub_cac, qsub_loc, run_local

# Global Variables
NODE_MEM=45000.0
NODE_CORES=12
SINGLE_CORE=1
MEM_PER_CORE= int(float(NODE_MEM) / NODE_CORES)
# wt=walltime
WT_SHORT= "10:00:00" #24.00.00
WT_LONG= "60:00:00" #"100:00:00"
WT_ELONG="100:00:00"


def bed_tools_isec(bedtools_path, fileA, fileB, isecfile, local=False):
    '''
    Print intersection between file A and file B
    Print only unique entries infile A vs file B 
    '''
    bed_cmd = os.path.join(bedtools_path,"intersectBed")
    args = [bed_cmd, '-a', fileA, '-b', fileB, '-wa', '-u', '>', isecfile]    
    if local:
        args = [bed_cmd, '-a', fileA, '-b', fileB, '-wa', '-u']
        args = (args,isecfile)
    
    return args

def bed_tools_complement(bedtools_path, fileA, fileB, isecfile, local):
    '''
    Print intersection complement between fileA and fileB
    '''
    bed_cmd = os.path.join(bedtools_path,"intersectBed")
    
    args = [bed_cmd, '-a', fileA, '-b', fileB, '-wa', '-v', '>', isecfile]    
    
    if local:
        args = [bed_cmd, '-a', fileA, '-b', fileB, '-wa', '-v']
        args = (args,isecfile)

    return args


def isec_file_name(descriptorA, descriptorB, compl=False):
    '''
    '''
    if not compl:
        fname='%s_%s.isec'%(descriptorA,descriptorB)
    else:
        fname='%s_%s.compl'%(descriptorA,descriptorB)
    return fname
    
    
def venn(fileA, fileB, descriptorA, descriptorB, bedtools_path, folder, jobrunfunc, local):
    '''
    Write a file with venn diagram values
    '''
    
    isec_name = isec_file_name(descriptorA, descriptorB, compl=False)
    compl_name = isec_file_name(descriptorA, descriptorB, compl=True)
    
    # command for intersection of files
    isec_args = bed_tools_isec(bedtools_path, fileA, fileB, os.path.join(folder,isec_name),local) 
    compl_args = bed_tools_complement(bedtools_path, fileA, fileB, os.path.join(folder,compl_name), local)    
    
    if not local:
        cmd_args=[",".join(isec_args).replace(',',' '),
                  ",".join(compl_args).replace(',',' ')]
        cmd_args=",".join(cmd_args).replace(',','\n')
    else:
        cmd_args=[isec_args,compl_args]
    
    return cmd_args, os.path.join(folder,isec_name), os.path.join(folder,compl_name)
    

def main_venn(fileA, fileB, whoisA, whoisB, bedtools_path, folder, jobrunfunc, local):
    '''
    '''
    my_email='alebalbin@gmail.com'
    jobidvcf=None
    jobn='bedtools'
    
    cmd_args_AvsB, isecAB, complAB = venn(fileA, fileB, whoisA, whoisB, bedtools_path, folder, jobrunfunc,local)
    cmd_args_BvsA, isecBA, complBA = venn(fileB, fileA, whoisB, whoisA, bedtools_path, folder, jobrunfunc, local)
    
    if not local:
        cmd = [cmd_args_AvsB, cmd_args_BvsA]
        cmd = ",".join(cmd).replace(',','\n')
        jobidvcf = jobrunfunc(jobn+'isec', cmd, SINGLE_CORE, cwd=None, walltime=WT_SHORT, pmem=None, 
                               deps=jobidvcf, stdout=None, email_addresses=my_email)
    else:

        cmds = cmd_args_AvsB+cmd_args_BvsA
        for cmd in cmds:
            jobidvcf = jobrunfunc(jobn+'isec', cmd[0], SINGLE_CORE, cwd=None, walltime=WT_SHORT, pmem=None, 
                               deps=jobidvcf, stdout=cmd[1], email_addresses=my_email)
    
    isec_dict={isecAB:[whoisA, whoisB,'isec'], complAB:[whoisA, whoisB,'compl'],
               isecBA:[whoisB, whoisA,'isec'], complBA:[whoisB, whoisA,'compl']}
    
    return jobidvcf, isec_dict


def count_lines_file(ifile):
    '''
    '''
    p2 = subprocess.Popen(['wc','-l',ifile],stdout=subprocess.PIPE)
    nlines = p2.stdout.readline().split(' ')[0]
    
    return nlines 
    
def make_venn_table(isec_files,oufile):
    '''
    write a venn square data
    '''
    oufile = open(oufile,'w')
    header=['sets_compared','comparison','#snvs']
    oufile.write(",".join(header).replace(',','\t')+'\n')
    for ifile, desc in isec_files.iteritems():
        # basically counting the number of lines
        nsnvs = count_lines_file(ifile)
        oline = [desc[0], desc[1], desc[2], nsnvs] 
        oufile.write(",".join(oline).replace(',','\t')+'\n')


    
if __name__ == '__main__':
    
    optionparser = OptionParser("usage: %prog [options] ")
    # Specific options
    optionparser.add_option("-b", "--bed_tools_path", dest="bed_tools_path",
                            help="bed_tools_path")

    optionparser.add_option("-f", "--query_vcf_file", dest="query_vcf_file",
                            help="query_vcf_file")
    optionparser.add_option("-o", "--others_vcf_files",dest="others_vcf_files",
                            help="others_vcf_files. Pass vcf files without .gz extention even if the file .gz is there")
    optionparser.add_option("-A", "--whoisA", dest="whoisA",
                            help="whoisA")
    optionparser.add_option("-B", "--whoisB", dest="whoisB",
                            help="whoisB")
    optionparser.add_option("-p", "--outfolder", dest="outfolder",
                            help="pathfolder")


    optionparser.add_option("--local", dest="local", action="store_true", default=False)
    optionparser.add_option("--cluster", dest="cluster", action="store_true", default=False)

    (options, args) = optionparser.parse_args()
    print options
    
    logfile = open(os.path.join(options.outfolder,'README'+options.whoisA+'_'+options.whoisB+'.log'),'w')
    logfile.write('Run call:\n')
    [logfile.write(",".join([str(opt),str(val)]).replace(',','\t')+'\n') for opt, val in options.__dict__.iteritems()]

    
    if not (options.local ^ options.cluster):
        print "Must set either --local or --cluster to run job"
    if options.local:
        jobrunfunc = run_local
    elif options.cluster:
        jobrunfunc = qsub
    
    if (options.whoisA == '' or options.whoisB == ''):
        whoisA, whoisB = 'fileA', 'fileB'
    else:
        whoisA, whoisB = options.whoisA, options.whoisB
    
    
    jobidvcf, isec_files_dict = main_venn(options.query_vcf_file, options.others_vcf_files, whoisA, whoisB, 
                                    options.bed_tools_path, options.outfolder, jobrunfunc, options.local)
    
    isec_files_dict[options.query_vcf_file]=[whoisA,whoisA,'total']
    isec_files_dict[options.others_vcf_files]=[whoisB,whoisB,'total']
    oufile=os.path.join(options.outfolder,"venn."+options.whoisA+'_'+options.whoisB+".stats")
    make_venn_table(isec_files_dict, oufile)




    